#include "WProgram.h"
// Ted Carancho's MikroQuad - November 2008
// An Arduino based quadrocopter using the Sparkfun
// 5DOF IMU and IDG300 Dual Axis Gyro
// This version will be able to use gyros for stability (acrobatic mode)
// or accelerometers (stable mode).

#include <EEPROM.h>
#include <ServoTimer2.h>

// ******************** Initialize Variables ********************
#define BAUD 38400

// Define sensor pin assignments
#define XACCELPIN 0
#define YACCELPIN 1
#define ZACCELPIN 2
#define PITCHRATEPIN 3
#define ROLLRATEPIN 4
#define YAWRATEPIN 5

// Define EEPROM storage
#define RATE_PGAIN_ADR 0
#define RATE_IGAIN_ADR 4
#define RATE_DGAIN_ADR 8
#define POS_PGAIN_ADR 12
#define POS_IGAIN_ADR 16
#define POS_DGAIN_ADR 20
#define XMITFACTOR_ADR 24


// Define transmitter commands
#define TIMEOUT 2500
#define THROTTLEPIN 4
#define ROLLPIN 2
#define PITCHPIN 3
#define YAWPIN 6
#define MODEPIN 7
#define AUXPIN 5
void setup();
void loop ();
float readFloatSerial();
void commandAllMotors(int motorCommand);
float readFloat(int address);
void writeFloat(float value, int address);
void interruptHandler();
void comma();
void initGyro1DKalman(struct Gyro1DKalman *filterdata, float Q_angle, float Q_gyro, float R_angle);
void predictKalman(struct Gyro1DKalman *filterdata, const float dotAngle, const float dt);
float updateKalman(struct Gyro1DKalman *filterdata, const float angle_m);
float updatePID(float targetPosition, float currentPosition, float *lastPosition, float *iState);
int smooth(int rawIn, int *smooth_array);
int findZero(int channel);
void zeroGyros();
void zeroAccelerometers();
volatile int throttle = 0;
volatile int roll = 0;
volatile int pitch = 0;
volatile int yaw = 0;
volatile int mode = 0; // >1500 stable mode,  <1500 acrobatic mode
volatile int aux= 0;
int rollCommand = 1500;
int pitchCommand = 1500;
int yawCommand = 1500;
int rollCommandZero, pitchCommandZero, yawCommandZero;
// Controls the strength of the commands sent from the transmitter
// xmitFactor ranges from 0.01 - 1.0 (0.01 = weakest, 1.0 - strongest)
float xmitFactor;


// Define motor commands
#define MINCOMMAND 1000
#define MAXCOMMAND 2000
#define FRONTMOTORPIN 8
#define REARMOTORPIN 9
#define RIGHTMOTORPIN 10
#define LEFTMOTORPIN 11
ServoTimer2 frontMotor;
ServoTimer2 rearMotor;
ServoTimer2 rightMotor;
ServoTimer2 leftMotor;
int frontCommand;
int rearCommand;
int rightCommand;
int leftCommand;
int rollMotorCommand = 1500;
int pitchMotorCommand = 1500;
int yawMotorCommand = 1500;

// Define sensor data
#define ZEROLIMIT 2
// Accel setup
int rollSensor, pitchSensor, yawSensor;
int xAccelZero, yAccelZero, zAccelZero;
int zAccSensor;
float rollLevelAdjust, pitchLevelAdjust;
float levelFactor = 1.0;

// Gyro setup
int rollGyroZero, pitchGyroZero, yawGyroZero;
// rate = (N-512)/1024 * 3.3V * (1 deg/s)/.002V * (Pi radians)/(180 deg)
// rate = (N-512)/1024 * (double) 28.783, where N is measured A/D value
float convert2Radians = 1 / 29473.792;

// If AREF = 3.3V, then A/D is 931 at 3V and 465 = 1.5V
// Scale gyro output (-465 to +465) to motor commands (1000 to 2000)
// use y = mx + b
float mMotorRate = 1.0753; // m = (y2 - y1) / (x2 - x1) = (2000 - 1000) / (465 - (-465))
float bMotorRate = 1500;   // b = y1 - m * x1

// Scale accel output (-100 to +100) to motor commands (1000 to 2000)
float mMotorRoll = 5;
float bMotorRoll = 1500;

// Scale accel output (+100 to -100) to motor commands (1000 to 2000)
float mMotorPitch = -2.5;
float bMotorPitch = 1500;

// Smoothing parameters
int rollSmoothArray [16];
int pitchSmoothArray [16];
int yawSmoothArray [16];

// PID Values
#define WINDUP_GUARD_GAIN 100.0
float pTerm, iTerm, dTerm; 
float pgain = 0;
float igain = 0;
float dgain = 0;
float rate_pgain = 0;
float rate_igain = 0;
float rate_dgain = 0;
float pos_pgain = 0;
float pos_igain = 0;
float pos_dgain = 0;
float iRollState = 0;
float lastRollPosition = 0;
float iPitchState = 0;
float lastPitchPosition = 0;
float iYawState = 0;
float lastYawPosition = 0;

// Kalman filter
struct Gyro1DKalman {
  float x_angle, x_bias;
  float P_00, P_01, P_10, P_11;	
  float Q_angle, Q_gyro;
  float R_angle;
};
struct Gyro1DKalman rollFilter;
struct Gyro1DKalman pitchFilter;
float dt;

// Communication
volatile char queryType = 'X';
char string[32];
byte armed;

// Interrupt Handler
volatile byte timeSlot;

// Timing
long previousTime = 0;
long currentTime = 0;
long deltaTime = 0;

// ******************** Setup MikroQuadAero ********************
void setup() {
  Serial.begin(BAUD);
  analogReference(EXTERNAL); // Current external ref is connected to 3.3V
  pinMode(THROTTLEPIN, INPUT);
  pinMode(ROLLPIN, INPUT);
  pinMode(PITCHPIN, INPUT);
  pinMode(YAWPIN, INPUT);
  pinMode(MODEPIN, INPUT);
  pinMode(AUXPIN, INPUT);
  
  // Configure motors
  frontMotor.attach(FRONTMOTORPIN);
  rearMotor.attach(REARMOTORPIN);
  rightMotor.attach(RIGHTMOTORPIN);
  leftMotor.attach(LEFTMOTORPIN);
  commandAllMotors(MINCOMMAND);

  // Read PID values form EEPROM
  rate_pgain = readFloat(RATE_PGAIN_ADR);
  rate_igain = readFloat(RATE_IGAIN_ADR);
  rate_dgain = readFloat(RATE_DGAIN_ADR);
  pos_pgain = readFloat(POS_PGAIN_ADR);
  pos_igain = readFloat(POS_IGAIN_ADR);
  pos_dgain = readFloat(POS_DGAIN_ADR);

  // Read transmit factor from EEPROM
  xmitFactor = readFloat(XMITFACTOR_ADR);

  // Setup interrupt to trigger on pin D2 (roll receiver pin)
  attachInterrupt(0, interruptHandler, RISING);
  armed = 0;
  
  // Wait for ESC's to initialize
  delay(3000);

  // Calibrate sensors
  zeroGyros();
  zeroAccelerometers();
  initGyro1DKalman(&rollFilter, 0.0001, 0.0003, 0.02);
  initGyro1DKalman(&pitchFilter, 0.0001, 0.0003, 0.02);

  // Continue to wait for ESC's to initialize
  delay(2000);

  previousTime = millis();
}

// ******************** Main MikroQuadAero Loop ********************
void loop () {
  // Measure loop rate
  currentTime = millis();
  deltaTime = currentTime - previousTime;
  previousTime = currentTime;
  
  // Send configuration commands from transmitter
  if (throttle < 1010) {
    // Disarm motors (throttle down, yaw left)
    if (yaw < (MINCOMMAND + 100) && armed == 1) {
      armed = 0;
      commandAllMotors(MINCOMMAND);
    }
    
    // Zero sensors (throttle down, yaw left, roll left)
    if (yaw < (MINCOMMAND + 100) && (roll > (MAXCOMMAND - 100))) {
      zeroGyros();
      zeroAccelerometers();
    }
    
    // Arm motors (throttle down, yaw right)  
    if (yaw > (MAXCOMMAND - 100) && armed == 0) {
      armed = 1;
      commandAllMotors(MINCOMMAND + 75);
      delay(500);
      commandAllMotors(MINCOMMAND);
    }
  }

  // Read Sensors
  // Average sensor values and center around zero
  rollSensor = smooth(analogRead(ROLLRATEPIN), rollSmoothArray) - rollGyroZero;
  pitchSensor = smooth(analogRead(PITCHRATEPIN), pitchSmoothArray) - pitchGyroZero;
  yawSensor = smooth(analogRead(YAWRATEPIN), yawSmoothArray) - yawGyroZero;

  if (mode < 1500) {
    // Acrobatic Mode
    rollLevelAdjust = 0;
    pitchLevelAdjust = 0;
  }
  else {
    // Stable Mode
    zAccSensor = analogRead(ZACCELPIN) + zAccelZero;
    dt = deltaTime / 1000.0;

    predictKalman(&rollFilter, rollSensor * convert2Radians, dt);
    rollLevelAdjust = updateKalman(&rollFilter, atan2(zAccSensor, analogRead(YACCELPIN) + yAccelZero)) * levelFactor;

    predictKalman(&pitchFilter, pitchSensor * convert2Radians, dt);
    pitchLevelAdjust = updateKalman(&pitchFilter, atan2(zAccSensor, analogRead(XACCELPIN) + xAccelZero)) * levelFactor;
  }
  
  // Calculate Commanded Angles
  // Reduce transmitter commands using xmitFactor and center around 1500
  rollCommand = ((roll - 1500) * xmitFactor) + 1500 + rollLevelAdjust;
  pitchCommand = ((pitch - 1500) * xmitFactor) + 1500 + pitchLevelAdjust;
  yawCommand = ((yaw - 1500) * xmitFactor) + 1500;
  
  // Update PID
  // Also convert sensor values to same range as PWM
  rollMotorCommand = updatePID((mMotorRate * rollSensor) + bMotorRate, rollCommand, &lastRollPosition, &iRollState);
  pitchMotorCommand = updatePID((mMotorRate * pitchSensor) + bMotorRate, pitchCommand, &lastPitchPosition, &iPitchState);
  yawMotorCommand = updatePID((mMotorRate * yawSensor) + bMotorRate, yawCommand, &lastYawPosition, &iYawState);   

  // Calculate motor commands
  frontCommand = constrain((throttle - pitchMotorCommand + yawMotorCommand), MINCOMMAND, MAXCOMMAND);
  rearCommand = constrain((throttle + pitchMotorCommand + yawMotorCommand), MINCOMMAND, MAXCOMMAND);
  rightCommand = constrain(throttle - rollMotorCommand - yawMotorCommand, MINCOMMAND, MAXCOMMAND);
  leftCommand = constrain((throttle + rollMotorCommand - yawMotorCommand), MINCOMMAND, MAXCOMMAND);
    
  // Command motors
  if (armed) {
    frontMotor.write(frontCommand);
    rearMotor.write(rearCommand);
    rightMotor.write(rightCommand);
    leftMotor.write(leftCommand);
  }
  
  // Check for serial message
  if (Serial.available()) {
    queryType = Serial.read();
    if (queryType == 'U') {
      // Update PID values    
      rate_pgain = readFloatSerial();
      writeFloat(rate_pgain, RATE_PGAIN_ADR);
      rate_igain = readFloatSerial();
      writeFloat(rate_igain, RATE_IGAIN_ADR);
      rate_dgain = readFloatSerial();
      writeFloat(rate_dgain, RATE_DGAIN_ADR);
      pos_pgain = readFloatSerial();
      writeFloat(pos_pgain, POS_PGAIN_ADR);
      pos_igain = readFloatSerial();
      writeFloat(pos_igain, POS_IGAIN_ADR);
      pos_dgain = readFloatSerial();
      writeFloat(pos_dgain, POS_DGAIN_ADR);
      // Update transmitter factor
      xmitFactor = readFloatSerial();
      writeFloat(xmitFactor, XMITFACTOR_ADR);
      queryType = 'X';
    }
  }
}

// ******************** Subroutines ********************
// Used to read floating point values from the serial port
float readFloatSerial() {
  int i = -1;
  char data[128] = "";
  
  do {
    while (Serial.available() == 0) delay(50); // wait if no data
    i++;
    data[i] = Serial.read();
  } while (data[i] != ';'); // read serial port until semi-colon found
  
  return strtod(data, NULL);
}

// Sends commands to all motors
void commandAllMotors(int motorCommand) {
  frontMotor.write(motorCommand);
  rearMotor.write(motorCommand);
  rightMotor.write(motorCommand);
  leftMotor.write(motorCommand);
}

// Utilities for writing and reading from the EEPROM

float readFloat(int address) {
  union floatStore {
    byte floatByte[4];
    float floatVal;
  } floatOut;
  
  for (int i = 0; i < 4; i++) 
    floatOut.floatByte[i] = EEPROM.read(address + i);
  return floatOut.floatVal;
}

void writeFloat(float value, int address) {
  union floatStore {
    byte floatByte[4];
    float floatVal;
  } floatIn;
  
  floatIn.floatVal = value;
  for (int i = 0; i < 4; i++) 
    EEPROM.write(address + i, floatIn.floatByte[i]);
}

// Interrupts are generated every 20ms because
// Roll pin from transmitter triggers interrupt (50MHz, PWM)
// Modulo division by 5 causes each time slot to happen every 100ms

void interruptHandler() {
  switch ((timeSlot += 1) % 5) {
  case 0:
    // Read Receiver
    // This has been  tested with AR6100 and AR6200 Spektrum receivers
    roll = pulseIn(ROLLPIN, HIGH, TIMEOUT);
    aux = pulseIn(AUXPIN, HIGH, TIMEOUT);
    mode = pulseIn(MODEPIN, HIGH, TIMEOUT);
    pitch = pulseIn(PITCHPIN, HIGH, TIMEOUT);
    throttle = pulseIn(THROTTLEPIN, HIGH, TIMEOUT);
    yaw = pulseIn(YAWPIN, HIGH, TIMEOUT);
    break;
  case 2:
    // Send Data
    switch (queryType) {
    case 'X': // send no debug messages
      break;
    case 'A': // send all data
      Serial.print(deltaTime);
      comma();
      Serial.print(rollSensor);
      comma();
      Serial.print(pitchSensor);
      comma();
      Serial.print(yawSensor);
      comma();
      Serial.print(throttle);
      comma();
      Serial.print(rollMotorCommand);
      comma();
      Serial.print(pitchMotorCommand);
      comma();
      Serial.print(yawMotorCommand);
      comma();
      Serial.print(frontCommand);
      comma();
      Serial.print(rearCommand);
      comma();
      Serial.print(rightCommand);
      comma();
      Serial.print(leftCommand);
      comma();
      Serial.print(armed, BIN);
      comma();
      Serial.println(mode);
      break;
    case 'S': // send sensor data
      Serial.print(rollSensor);
      comma();
      Serial.print(pitchSensor);
      comma();
      Serial.print(yawSensor);
      comma();
      Serial.print(dtostrf(rollLevelAdjust, 1, 3, string));
      comma();
      Serial.println(dtostrf(pitchLevelAdjust, 1, 3, string));
      break;
    case 'D': // raw read sensor data
      Serial.print(analogRead(ROLLRATEPIN));
      comma();
      Serial.print(analogRead(PITCHRATEPIN));
      comma();
      Serial.print(analogRead(YAWRATEPIN));
      comma();
      Serial.print(analogRead(XACCELPIN));
      comma();
      Serial.print(analogRead(YACCELPIN));
      comma();
      Serial.println(analogRead(ZACCELPIN));
      break;
    case 'P': // send PID values
      Serial.print(dtostrf(rate_pgain, 1, 2, string));
      comma();
      Serial.print(dtostrf(rate_igain, 1, 2, string));
      comma();
      Serial.print(dtostrf(rate_dgain, 1, 2, string));
      comma();
      Serial.print(dtostrf(pos_pgain, 1, 2, string));
      comma();
      Serial.print(dtostrf(pos_igain, 1, 2, string));
      comma();
      Serial.print(dtostrf(pos_dgain, 1, 2, string));
      comma();
      Serial.println(dtostrf(xmitFactor, 1, 2, string));
      queryType = 'X';
      break;
    case 'T': // read processed transmitter values
      Serial.print(dtostrf(xmitFactor, 1, 2, string));
      comma();
      Serial.print(rollCommand);
      comma();
      Serial.print(pitchCommand);
      comma();
      Serial.println(yawCommand);
      break;
    case 'R': // send receiver values
      Serial.print(throttle);
      comma();
      Serial.print(roll);
      comma();
      Serial.print(pitch);
      comma();
      Serial.print(yaw);
      comma();
      Serial.print(mode);
      comma();
      Serial.println(aux);
      break;
    }
  case 3: // check flight mode and adjust PID values to use
    if (mode < 1500) {
      // acrobatic mode
      pgain = rate_pgain;
      igain = rate_igain;
      dgain = rate_dgain;
    }
    else {
      // stable mode
      pgain = pos_pgain;
      igain = pos_igain;
      dgain = pos_dgain;
    }
  default:
  break;
  }
}

void comma() {
  Serial.print(',');
}

// The sensor platform used is the SparkFun 5DOF IMU found at:
// http://www.sparkfun.com/commerce/product_info.php?products_id=741
// The Kalman filter implementation is directly taken from the work
// of Tom Pycke at: http://tom.pycke.be/mav/90/sparkfuns-5dof
// The AREF pin is connected to 3.3V to better use the 10 bit A/D

void initGyro1DKalman(struct Gyro1DKalman *filterdata, float Q_angle, float Q_gyro, float R_angle) {
	filterdata->Q_angle = Q_angle;
	filterdata->Q_gyro  = Q_gyro;
	filterdata->R_angle = R_angle;
}

void predictKalman(struct Gyro1DKalman *filterdata, const float dotAngle, const float dt) {
	filterdata->x_angle += dt * (dotAngle - filterdata->x_bias);
	filterdata->P_00 +=  - dt * (filterdata->P_10 + filterdata->P_01) + filterdata->Q_angle * dt;
	filterdata->P_01 +=  - dt * filterdata->P_11;
	filterdata->P_10 +=  - dt * filterdata->P_11;
	filterdata->P_11 +=  + filterdata->Q_gyro * dt;
}

float updateKalman(struct Gyro1DKalman *filterdata, const float angle_m) {
	const float y = angle_m - filterdata->x_angle;
	const float S = filterdata->P_00 + filterdata->R_angle;
	const float K_0 = filterdata->P_00 / S;
	const float K_1 = filterdata->P_10 / S;
	
	filterdata->x_angle +=  K_0 * y;
	filterdata->x_bias  +=  K_1 * y;
	filterdata->P_00 -= K_0 * filterdata->P_00;
	filterdata->P_01 -= K_0 * filterdata->P_01;
	filterdata->P_10 -= K_1 * filterdata->P_00;
	filterdata->P_11 -= K_1 * filterdata->P_01;

	return filterdata->x_angle;
}

// Taken from http://www.arduino.cc/playground/Main/BarebonesPIDForEspresso

float updatePID(float targetPosition, float currentPosition, float *lastPosition, float *iState)
{
  // these local variables can be factored out if memory is an issue, 
  // but they make it more readable
  float error;
  float windupGaurd;

  // determine how badly we are doing
  error = targetPosition - currentPosition;

  // the pTerm is the view from now, the pgain judges 
  // how much we care about error we are this instant.
  pTerm = pgain * error;

  // iState keeps changing over time; it's 
  // overall "performance" over time, or accumulated error
  *iState += error;

  // to prevent the iTerm getting huge despite lots of 
  //  error, we use a "windup guard" 
  // not necessary, but this makes windup guard values 
  // relative to the current iGain
  windupGaurd = WINDUP_GUARD_GAIN / igain;  

  if (*iState > windupGaurd) 
    *iState = windupGaurd;
  else if (*iState < -windupGaurd) 
    *iState = -windupGaurd;
  iTerm = igain * (*iState);

  // the dTerm, the difference between the current
  //  and last reading, indicated the "speed," 
  // how quickly the reading is changing. (aka. Differential)
  dTerm = (dgain * (currentPosition - *lastPosition));
  *lastPosition = currentPosition;
  return  pTerm + iTerm - dTerm;
}

// Utilities for sensor measurements

// Removes noise from sensor data
int smooth(int rawIn, int *smooth_array) {
  long total = 0;
  static int i;

  i = (i + 1) % 16;
  smooth_array[i] = rawIn;
  for (int j = 0; j < 16; j++) total += smooth_array[j];
  return total >> 4 ;
}

// Used to average the sensors measurements when everything is level
int findZero(int channel) {
  long zero = 0;
  
  for (int i=0; i< 32; i++) zero += analogRead(channel);
  return zero >> 5;
}

// Allows user to zero gyros on command
void zeroGyros() {
  do {
    rollGyroZero = findZero(ROLLRATEPIN);
  } while (abs(analogRead(ROLLRATEPIN) - rollGyroZero) > ZEROLIMIT);
  
  do {
    pitchGyroZero = findZero(PITCHRATEPIN);
  } while (abs(analogRead(PITCHRATEPIN) - pitchGyroZero) > ZEROLIMIT);

  do {
    yawGyroZero = findZero(YAWRATEPIN);
  } while (abs(analogRead(YAWRATEPIN) - yawGyroZero) > ZEROLIMIT);
}

// Allows user to zero accelerometers on command
void zeroAccelerometers() {
  do {
    xAccelZero = findZero(XACCELPIN);
  } while (abs(analogRead(XACCELPIN) - xAccelZero) > ZEROLIMIT);
  
  do {
    yAccelZero = findZero(YACCELPIN);
  } while (abs(analogRead(YACCELPIN) - yAccelZero) > ZEROLIMIT);
  
  do {
    zAccelZero = findZero(ZACCELPIN);
  } while (abs(analogRead(ZACCELPIN) - zAccelZero) > ZEROLIMIT);
}
int main(void)
{
	init();

	setup();
    
	for (;;)
		loop();
        
	return 0;
}

